import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import KFold
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics import r2_score
import matplotlib.pyplot as plt
import random

# 固定随机种子
def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)

set_seed(42)  # 固定随机种子

# 读取处理后的数据
df = pd.read_csv('processed_data_with_bands_all9.18.csv')
df = df[df['fy_clt'] == 2]

# 提取特征和目标
features = ['fy_cth_log', 'fy_ctt_scaled', 'fy_ctp_scaled',
            'fy_lat_sin', 'fy_lat_cos', 'fy_lon_sin', 'fy_lon_cos',
            'fy_band1_scaled', 'fy_band2_scaled', 'fy_band5_scaled', 'fy_band6_scaled', 'fy_band7_scaled',
            'fy_band12_scaled']
target = 'cloudsat_cbh'

X = df[features].values
y = df[target].values

# 对目标值进行标准化
scaler_y = StandardScaler()
y = scaler_y.fit_transform(y.reshape(-1, 1)).flatten()

# 定义 KFold 进行交叉验证
kfold = KFold(n_splits=5, shuffle=True, random_state=42)

# 定义 CNN-DNN 模型
# class CNN_DNN(nn.Module):
#     def __init__(self, input_size):
#         super(CNN_DNN, self).__init__()
#         self.conv1 = nn.Conv1d(1, 32, kernel_size=3, padding=1)
#         self.bn1 = nn.BatchNorm1d(32)
#         self.conv2 = nn.Conv1d(32, 64, kernel_size=3, padding=1)
#         self.bn2 = nn.BatchNorm1d(64)
#         self.conv3 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
#         self.bn3 = nn.BatchNorm1d(128)
#
#         self.fc1 = nn.Linear(128 * input_size, 256)
#         self.fc2 = nn.Linear(256, 128)
#         self.fc3 = nn.Linear(128, 1)
#
#         self.relu = nn.LeakyReLU(0.01)
#         self.dropout = nn.Dropout(p=0.4)
#
#     def forward(self, x):
#         x = self.relu(self.bn1(self.conv1(x)))
#         x = self.relu(self.bn2(self.conv2(x)))
#         x = self.relu(self.bn3(self.conv3(x)))
#         x = x.view(x.size(0), -1)
#         x = self.relu(self.fc1(x))
#         x = self.dropout(x)
#         x = self.relu(self.fc2(x))
#         x = self.fc3(x)
#         return x

# 模型更新
class OptimizedCNN_DNN(nn.Module):
    def __init__(self, input_size):
        super(OptimizedCNN_DNN, self).__init__()
        # 增加卷积层和残差连接
        self.conv1 = nn.Conv1d(1, 64, kernel_size=3, padding=1)  # 增加通道数
        self.bn1 = nn.BatchNorm1d(64)
        self.conv2 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm1d(128)

        # 新增卷积层
        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm1d(256)

        # 全连接层
        self.fc1 = nn.Linear(256 * input_size, 512)  # 增加全连接层的神经元数
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 128)
        self.fc4 = nn.Linear(128, 1)  # 输出层

        # Dropout 层，避免过拟合
        self.relu = nn.LeakyReLU(0.01)
        self.dropout = nn.Dropout(p=0.3)  # 调低 Dropout 概率

    def forward(self, x):
        # 前向传播过程
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.relu(self.bn3(self.conv3(x)))  # 新增的卷积层
        x = x.view(x.size(0), -1)  # 平展为全连接层的输入
        x = self.relu(self.fc1(x))
        x = self.dropout(x)  # 增加 Dropout
        x = self.relu(self.fc2(x))
        x = self.dropout(x)  # 再次 Dropout
        x = self.relu(self.fc3(x))
        x = self.fc4(x)
        return x

# 记录每个 fold 的 R²
r2_scores = []

# 进行 K 折交叉验证
for fold, (train_idx, val_idx) in enumerate(kfold.split(X, y)):
    print(f"\nFold {fold + 1}")

    # 划分训练集和验证集
    X_train, X_val = X[train_idx], X[val_idx]
    y_train, y_val = y[train_idx], y[val_idx]

    # 转换为 PyTorch 张量
    X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
    y_train_tensor = torch.tensor(y_train, dtype=torch.float32).view(-1, 1)
    X_val_tensor = torch.tensor(X_val, dtype=torch.float32)
    y_val_tensor = torch.tensor(y_val, dtype=torch.float32).view(-1, 1)

    # 创建数据加载器
    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
    val_dataset = TensorDataset(X_val_tensor, y_val_tensor)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

    # 使用 GPU
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # model = CNN_DNN(X_train.shape[1]).to(device)
    model = OptimizedCNN_DNN(X_train.shape[1]).to(device)

    # 使用 MSE 作为损失函数
    criterion = nn.MSELoss()

    # 定义优化器和动态学习率调整
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5, verbose=True)

    # 训练模型
    num_epochs = 200
    early_stopping_patience = 20
    early_stopping_counter = 0
    best_val_loss = float('inf')

    for epoch in range(num_epochs):
        model.train()
        for X_batch, y_batch in train_loader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            optimizer.zero_grad()
            outputs = model(X_batch.unsqueeze(1))
            loss = criterion(outputs, y_batch)
            loss.backward()
            optimizer.step()

        # 验证模型
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for X_batch, y_batch in val_loader:
                X_batch, y_batch = X_batch.to(device), y_batch.to(device)
                outputs = model(X_batch.unsqueeze(1))
                loss = criterion(outputs, y_batch)
                val_loss += loss.item()

        val_loss /= len(val_loader)
        print(f'Epoch [{epoch + 1}/{num_epochs}], Fold {fold + 1}, Validation Loss: {val_loss:.4f}')

        scheduler.step(val_loss)

        # 早停机制
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            early_stopping_counter = 0
        else:
            early_stopping_counter += 1
            if early_stopping_counter >= early_stopping_patience:
                print(f'Early stopping at epoch {epoch + 1}')
                break

    with torch.no_grad():
        # 将验证集张量移动到与模型相同的设备
        X_val_tensor = X_val_tensor.to(device)

        # 进行预测并将结果转换回 CPU 上以计算
        y_pred = model(X_val_tensor.unsqueeze(1)).cpu().numpy()  # 模型预测
        y_val_true = y_val_tensor.cpu().numpy()  # 确保目标值在 CPU 上

    # 反标准化目标值
    y_pred_true = scaler_y.inverse_transform(y_pred)
    y_val_true = scaler_y.inverse_transform(y_val_true)

    r2 = r2_score(y_val_true, y_pred_true)
    print(f'Fold {fold + 1} R² Score: {r2:.4f}')
    r2_scores.append(r2)

    # 计算误差并绘制频率分布图
    errors = y_val_true - y_pred_true
    plt.figure(figsize=(10, 6))
    plt.hist(errors, bins=50, color='blue', edgecolor='black', alpha=0.7)
    plt.title(f'Error Frequency Distribution (True - Predicted-group2-Modle) - Fold {fold + 1}')
    plt.xlabel('Error')
    plt.ylabel('Frequency')
    plt.show()

# 检查所有 R² 分数
print(f"R² scores for all folds: {r2_scores}")

# 检查 r2_scores 列表是否为空
if r2_scores:
    # 计算平均 R²，过滤掉无效值
    r2_scores_filtered = [r for r in r2_scores if not np.isnan(r)]
    if r2_scores_filtered:
        avg_r2 = np.mean(r2_scores_filtered)
        print(f'Average R² Score across {len(r2_scores_filtered)} valid folds: {avg_r2:.4f}')
    else:
        print("No valid R² scores to average.")
else:
    print("No R² scores calculated.")


# 计算平均 R²
avg_r2 = np.mean(r2_scores)
print(f'Average R² Score across {kfold.n_splits} folds: {avg_r2:.4f}')
